
#include <opencv2/core.hpp>
#include <opencv2/core/ocl.hpp>
#include <opencv2/core/base.hpp>
#include <opencv2/core/utility.hpp>
#include <opencv2/imgproc.hpp>
#include "joint_bilateralfilter.h"
#include <opencv2/ximgproc.hpp>
#include <stdio.h>
#include <algorithm>
#include <map>
#include <climits>
#include <iostream>

#define CLIP_UIN8(v_s32) (((v_s32) & (~255)) ? (uint8_t)((-v_s32)>>31) : (uint8_t)(v_s32))

using namespace std;
using namespace cv;


#ifdef _MSC_VER
#   pragma warning(disable: 4512)
#endif

#ifndef SQR
#define SQR(a) ((a)*(a))
#endif


class JointBilateralFilter_8u : public ParallelLoopBody
{
    Mat &joint, &src;
    Mat &dst;
    int radius, maxk;
    float scaleIndex;
    int *spaceOfs;
    float *spaceWeights, *expLUT;
    Mat &mask;

public:

    JointBilateralFilter_8u(Mat& joint_, Mat& src_, Mat& dst_, int radius_,
        int maxk_, int *spaceOfs_, float *spaceWeights_, float *expLUT_, Mat &mask_)
        :
        joint(joint_), src(src_), dst(dst_), radius(radius_), maxk(maxk_),
        spaceOfs(spaceOfs_), spaceWeights(spaceWeights_), expLUT(expLUT_), mask(mask_)
    {

        CV_DbgAssert(joint.rows == src.rows && src.rows == dst.rows + 2 * radius);
        CV_DbgAssert(joint.cols == src.cols && src.cols == dst.cols + 2 * radius);
    }

    void operator () (const Range& range) const CV_OVERRIDE
    {
        uint8_t mask_val;
        
        for (int i =  range.start; i <  range.end; i++)
        {
            for (int j = radius; j < src.cols - radius; j++)
            {
                uint8_t *jointCenterPixPtr = joint.ptr<uint8_t>(i) + j;
                uint8_t *srcCenterPixPtr = src.ptr<uint8_t>(i) + j;
                mask_val = mask.at<uint8_t>(i, j);

                if (!mask_val){
                    dst.at<uint8_t>(i, j) = src.at<uint8_t>(i, j);
                    continue;
                }

                uint8_t jointPix0 = *jointCenterPixPtr;
                float srcSum = 0;
                float wSum = 0.0f;

                for (int k = 0; k < maxk; k++)
                {
                    uchar *jointPix = (jointCenterPixPtr + spaceOfs[k]);
                    int alpha = std::abs(jointPix0 - (int)jointPix[0]);
                    float weight = spaceWeights[k] * expLUT[alpha];
                    uchar srcPix = *(srcCenterPixPtr + spaceOfs[k]);

                    srcSum += weight*srcPix;
                     wSum += weight;
                }
                int32_t v_s32 = (int32_t)(srcSum / wSum);
                dst.at<uint8_t>(i, j) = CLIP_UIN8(v_s32);
            }
        }
    }
};

void cv_jointBilateralFilter_8u(Mat& joint, Mat& src, Mat& dst, int d, double sigmaColor, double sigmaSpace,  Mat & mask)
{
    int i, j;
    int radius;
    
    CV_DbgAssert(joint.depth() == CV_8U && src.depth() == CV_8U);

    (void)mask;

    
    if (d <= 0)
        radius = cvRound(sigmaSpace*1.5);
    else
        radius = d / 2;
    radius = std::max(radius, 1);

    d = 2 * radius + 1;
    
    int jCn = joint.channels();

    double gaussColorCoeff = -0.5 / (sigmaColor*sigmaColor);
    double gaussSpaceCoeff = -0.5 / (sigmaSpace*sigmaSpace);

    vector<float> expLUTv(jCn*256);
    float *expLUT = &expLUTv[0];

    for ( i = 0; i < (int)expLUTv.size(); i++)
    {
        expLUT[i] = (float)std::exp(i * i * gaussColorCoeff);
    }

    size_t srcElemStep = src.step / src.elemSize();
    size_t jElemStep = joint.step / joint.elemSize();
    CV_Assert(srcElemStep == jElemStep);

    vector<float> spaceWeightsv(d*d);
    vector<int> spaceOfsJointv(d*d);
    float *spaceWeights = &spaceWeightsv[0];
    int *spaceOfsJoint = &spaceOfsJointv[0];

    int maxk = 0;
    for ( i = -radius; i <= radius; i++)
    {
        for ( j = -radius; j <= radius; j++)
        {
            double r2 = i*i + j*j;
            if (r2 > SQR(radius))
                continue;

            spaceWeights[maxk] = (float)std::exp(r2 * gaussSpaceCoeff);
            spaceOfsJoint[maxk] = (int)(i*jElemStep + j);
            maxk++;
        }
    }

    Range range(radius, src.rows - radius);
    parallel_for_(range, JointBilateralFilter_8u(joint, src, dst, radius, maxk, spaceOfsJoint, spaceWeights, expLUT, mask));


    {
    
        int  row;
        uint8_t *pix_src, *pix_dst;
        
        for(row = 0; row < radius ; ++row){
            pix_src = src.data + row * src.step[0];
            pix_dst = dst.data + row * dst.step[0];
        
            memcpy(pix_dst, pix_src, src.cols);
        }
        
        
        for(; row < src.rows - radius ; ++row){
            //left 
            pix_src = src.data + row * src.step[0];
            pix_dst = dst.data + row * dst.step[0];
        
            memcpy(pix_dst, pix_src, radius);
        
        
            //right 
            pix_src = src.data + row * src.step[0] + src.cols - radius;
            pix_dst = dst.data + row * dst.step[0] + dst.cols - radius;
        
            memcpy(pix_dst, pix_src, radius);
        
        }
        
        for(row = src.rows - radius ; row < src.rows ; ++row){
            pix_src = src.data + row * src.step[0];
            pix_dst = dst.data + row * dst.step[0];
        
            memcpy(pix_dst, pix_src, src.cols);
        }
        
        
        
        for(row = src.rows - radius ; row < src.rows ; ++row){
            pix_src = src.data + row * src.step[0];
            pix_dst = dst.data + row * dst.step[0];
        
            memcpy(pix_dst, pix_src, src.cols);
        }    

    }

}

